"""
name: id3
author: robocky
create time: 2016-12-31
description: Use id3 method to create a decision tree
"""
from math import log2
from makeSet import DictSet
from node import Node


def entropy(lst: list, key: str):
    """Calculate the entropy of signal.
    Entropy = sum(-pi * log2(pi))
    """
    if len(lst) > 0 and key in lst[0]:
        lst_num = len(lst)
        attr_set = set(data[key] for data in lst)
        attr_dict = {attr: len([data for data in lst if data[key] == attr]) for attr in attr_set}
        # print(attr_dict)
        return sum(-num / lst_num * log2(num / lst_num) for num in attr_dict.values())
    else:
        print("Des not have the key %s." % key)
        return 0


def gain(s: DictSet, attr: str, key: str):
    """Calculate the gain for the attr
    Gain(S, A) = Entropy(S) - sum(p(Sv) * Entropy(Sv))

    :param s: DictSet data.
    :param attr: The attr to split s.
    :param key: The key to calculate entropy.
    """
    sv_dict = s.split_data(attr)
    sv_entropy_list = []
    for sv in sv_dict.values():
        p_sv = len(sv) / len(s)
        sv_entropy_list.append(p_sv * entropy(sv, key))
    return entropy(s, key) - sum(sv_entropy_list)


def id3(s: DictSet, node: Node):
    """Use ID3 to generate a node."""
    out_key = s.get_out_key()
    keys = s.get_keys()
    keys.remove(out_key)
    gains = {key: gain(s, key, out_key) for key in keys}
    attr = max(gains, key=lambda k: gains[k])
    node.set_name(attr)
    split_s = s.split_data(attr)
    for attr_val, sv in split_s.items():
        if entropy(sv, out_key) == 0:
            node.add_leaf(attr_val, sv[0][out_key])
        else:
            node_son = node.add_branch(attr_val)
            sv_ds = DictSet(*s.get_keys())
            sv_ds.extend(sv)
            id3(sv_ds, node_son)

if __name__ == '__main__':
    # Test function entropy.
    print("Test function entropy.")
    lst_to_use = [
        {'id': 1, 'val': 'A'},
        {'id': 2, 'val': 'B'},
        {'id': 3, 'val': 'A'},
        {'id': 4, 'val': 'B'}
    ]
    print(entropy(lst_to_use, 'val'))
    # Test function gain.
    # A or (B and C)
    print("Test function gain.\nA or (B and C)")
    or_data = [
        ['True', 'True', 'True', 'True'],
        ['True', 'True', 'False', 'True'],
        ['True', 'False', 'True', 'True'],
        ['True', 'False', 'False', 'True'],
        ['False', 'True', 'True', 'True'],
        ['False', 'True', 'False', 'False'],
        ['False', 'False', 'True', 'False'],
        ['False', 'False', 'False', 'False'],
    ]
    s_to_use = DictSet('A', 'B', 'C', 'result')
    for data_to_add in or_data:
        s_to_use.add_data(data_to_add)
    print(entropy(s_to_use, 'result'))
    print(gain(s_to_use, 'A', 'result'))
    print(gain(s_to_use, 'B', 'result'))
    print(gain(s_to_use, 'C', 'result'))
    # Test ID3.
    print("Test ID3.")
    root = Node()
    id3(s_to_use, root)
    root.print_node()
    # Test ID3 use data.
    print("Test ID3 use data.")
    with open('data.txt') as f:
        s_data = DictSet()
        s_data.create_from_file(f)
        root = Node()
        id3(s_data, root)
        root.print_node()
